package cloud.tianai.mate.captcha.interceptor.impl;

import cloud.tianai.captcha.common.AnyMap;
import cloud.tianai.captcha.common.response.ApiResponse;
import cloud.tianai.captcha.common.response.CodeDefinition;
import cloud.tianai.captcha.common.util.CaptchaTypeClassifier;
import cloud.tianai.captcha.interceptor.Context;
import cloud.tianai.captcha.validator.common.model.dto.ImageCaptchaTrack;
import cloud.tianai.captcha.interceptor.CaptchaInterceptor;
import cloud.tianai.mate.captcha.validator.common.util.TrackUtils;
import lombok.SneakyThrows;
import ml.dmlc.xgboost4j.java.Booster;
import ml.dmlc.xgboost4j.java.DMatrix;
import ml.dmlc.xgboost4j.java.XGBoost;

import java.io.InputStream;
import java.util.List;

/**
 * @Author: 天爱有情
 * @date 2023/1/29 15:54
 * @Description xgb预测模型
 */
public class XGBTrackCaptchaInterceptor implements CaptchaInterceptor {
    public static final CodeDefinition DEFINITION = new CodeDefinition(50003, "xgb check fail");
    Booster model;


    @Override
    public String getName() {
        return "xgboost_track_check";
    }

    @SneakyThrows
    public XGBTrackCaptchaInterceptor(InputStream stream) {
        model = XGBoost.loadModel(stream);
    }

    @Override
    @SneakyThrows
    public ApiResponse<?> afterValid(Context context, String type, ImageCaptchaTrack imageCaptchaTrack, AnyMap validData, ApiResponse<?> basicValid) {
        if (!basicValid.isSuccess()) {
            return context.getGroup().afterValid(context, type, imageCaptchaTrack, validData, basicValid);
        }
        Object preReturn = context.getPreReturnData();
        if (preReturn!= null && !((ApiResponse<?>) preReturn).isSuccess()) {
            return (ApiResponse<?>) preReturn;
        }
        if (!CaptchaTypeClassifier.isSliderCaptcha(type)) {
            // 不是滑动验证码的话暂时跳过
            return ApiResponse.ofSuccess();
        }
        List<Double> features = context.getData(TrackFeaturesGenerator.TRACK_FEATURES_KEY, List.class);
        if (features == null) {
            features = TrackUtils.features(imageCaptchaTrack.getTrackList());
            context.putData(TrackFeaturesGenerator.TRACK_FEATURES_KEY, features);
        }
        // 先进行一下基本校验
//        Double yMin = features.get(2);
//        Double yMax = features.get(3);
//        Double totalTime = features.get(4);
//        Double yStd = features.get(10);
//        Double tStd = features.get(12);
//        if (yMin >= yMax || totalTime < 500 || yStd == 0 || tStd == 0) {
//            return false;
//        }

        float[] array = new float[features.size() + 1];
        array[0] = 1;
        for (int i = 0; i < features.size(); i++) {
            array[i + 1] = features.get(i).floatValue();
        }
        DMatrix dMatrix = new DMatrix(array, 1, array.length, 0.0f);
        float[][] predicts = model.predict(dMatrix);
        float predict = predicts[0][0];
        boolean check = predict > 0.5;
        if (check) {
            return ApiResponse.ofSuccess();
        }
        context.end();
        return ApiResponse.ofMessage(DEFINITION);
    }
}
